import source
import os
import soundfile as sf
from ai.utils import utils_file
from ai import AiConstant
import threading
import tqdm
import project_config
from ai.utils import utils_file

log_path = utils_file.join_path(AiConstant.LOG_PATH, project_config.PROJECT_NAME, 'fangyan_2.log')
logger = AiConstant.AI_LOGGER(log_path)

"""
最后使用的标签列表, 音频都是分过段的
text_data: 'audio_file_name\blabel'
wav_data: 'audio_file_name\baudio_file_path'
"""
text_data = []
wav_data = []

input_path = "/home/work_nfs7/dialects"
output_path = "./dialects/accent_label"
utils_file.makedir_for_file_or_dir(output_path)
# utils_file.makedir_for_file_or_dir(output_audio_path)
# for i in range(1,5):
#     lines = [line.strip() for line in open(input_path+"/dialects"+str(i)+".tsv", 'r').readlines()]
#     for (index, line) in enumerate(lines):
#         if index == 0:
#             continue
#         if "上海" in line:
#             shanghai.append(line)
#         elif "晋语" in line:
#             jinyu.append(line)
#         elif "粤语" in line:
#             yueyu.append(line)


# sichuan = []
# yueyu = []
# shanghai = []
# jinyu = []
# dongbeig = []
# lanying = []
# languages = {
#     "上海": shanghai = [],
#     "晋语": jinyu = [],
#     "粤语": yueyu = [],
#     "东北官话": dongbeig = [],
#     "兰银官话": lanying = [],
#     "冀鲁官话": jilug = [],
#     "北京官话": beijingg = [],
#     "江淮官话": jianghuaig = [],
#     "西南官话": xinang = [],
#     "胶辽官话": jiaoliao = [],
#     "湘语": xiangyu = [],
#     "吴语": wuyu = [],
#     "徽语": huiyu = [],
#     "赣语": ganyu = [],
#     "客家": kejia = [],
#     "闽南": minnan = [],
#     "闽语": minyu = [],
#     "海南": hainan = [],
#     "平话": pinghua = [],
# }

"""
languages : 按照文件目录最后一级名字得到key, 后跟data列表:'audio_path\tsample_num'
"""
languages = {
    "上海": [],
    "晋语": [],
    "粤语": [],
    "东北官话": [],
    "兰银官话": [],
    "冀鲁官话": [],
    "北京官话": [],
    "江淮官话": [],
    "西南官话": [],
    "胶辽官话": [],
    "湘语": [],
    "吴语": [],
    "徽语": [],
    "赣语": [],
    "客家": [],
    "闽南": [],
    "闽语": [],
    "海南": [],
    "平话": [],
}


def cut_wav(input_file, output_file, start_sample, end_sample):
    """
    读取一个音频文件, 并在指定的采样点范围进行切割, 然后生成一个新的文件
    """
    # 使用SoundFile库打开输入文件
    with sf.SoundFile(input_file, 'r') as infile:
        # 读取输入文件的采样率和帧数
        sample_rate = infile.samplerate
        total_samples = len(infile)

        # 确保结束采样点不超过文件的总采样数
        if end_sample > total_samples:
            end_sample = total_samples

        # 读取指定采样点范围的数据
        infile.seek(start_sample)
        data = infile.read(end_sample - start_sample, dtype='int16')

    # 使用SoundFile库创建输出文件并写入数据
    with sf.SoundFile(output_file, 'w', samplerate=sample_rate, channels=1, format='WAV', subtype='PCM_16') as outfile:
        outfile.write(data)


def write_file(data, accent: str):
    """
    data: 一个列表, 每个元素是一个字符串, 字符串格式如下: audio_file_path  sample_num
    accent: 一个字符串
    将长音频切分成都是20秒的固定时段音频
    """
    sub_text_data = []
    sub_wav_data = []
    time = 0
    logger.info(f'开始处理方言:{accent},待处理长音频数量为:' + str(len(data)))
    # tqdm_object = tqdm.tqdm(enumerate(data))
    # tqdm_object.total = len(data)
    # tqdm_object.desc = f'方言:{accent}handling'
    for (index, d) in enumerate(data):
        parts = d.split('\t')  # 数据位置  采样点个数
        wav_path = parts[0]
        file_name = wav_path[wav_path.rfind('/') + 1:]
        file_name = file_name[:file_name.rfind('.')]
        # file_name = file_name.strip(".wav")
        file_path = wav_path[:wav_path.rfind('/') + 1]
        output_wav = wav_path
        t = int(parts[1]) / 16000
        """
        将数据切割成20秒每段的小数据, 多出的部分舍弃
        """
        if t > 20:
            """
            主要是检查是否有还没切分的片段
            """
            num_segments = int(t / 20)  # 计算需要切割成多少段整20s的音频
            segment_duration = 20  # 每段的时长固定为20秒
            start_time = 0
            for i in range(num_segments):
                if (start_time + segment_duration) > t:
                    end_time = t  # # 其实没必要, 不可能超过
                else:
                    end_time = start_time + segment_duration
                output_wav = utils_file.join_path(file_path, file_name + f"_{i}.wav")
                if not os.path.exists(output_wav):
                    """
                    如果这个文件已经处理过,就无需再处理
                    """
                    logger.info(f"未处理音频: {output_wav}")
                    start_time = end_time
                    continue
                    # os.system(f"sox {wav_path} {output_wav} trim {start_time} {end_time}")
                    # cut_wav(wav_path, output_wav, start_time * 16000, end_time * 16000)
                # continue
                sub_text_data.append(f"{file_name}_{i} {accent}")  # 写入新的文本标签
                sub_wav_data.append(f"{file_name}_{i} {output_wav}")  # 写入新的音频文件路径
                start_time = end_time
        else:
            sub_text_data.append(f"{file_name} {accent}")  # 写入新的文本标签
            sub_wav_data.append(f"{file_name} {output_wav}")  # 写入新的音频文件路径
        time += t
        """
        要求每个口音的总时长不超过1000小时
        """
        if time / 3600 > 1000:
            break
    print('方言: ', accent, ' time: ', time / 3600, ' h.')
    text_data.extend(sub_text_data)
    wav_data.extend(sub_wav_data)


def loading_languages():
    """
    languages的value初始化还是空列表, 该函数根据dialects1-5.tsv文件读取数据对齐进行装载
    languages中装载的只是未切分的长音频, 不包含切分好的短音频的
    """
    for i in range(1, 5):
        """
        lines: 'audio_path\tsample_num'
        """
        lines = [line.strip() for line in
                 open(utils_file.join_path(input_path, "/dialects" + str(i) + ".tsv"), 'r', encoding='utf-8').readlines()]
        for (index, line) in enumerate(lines):
            if index == 0:
                continue
            for language in languages.keys():
                if language in line:
                    languages[language].append(line)
                    break


def mul_process(data_to_process, label):
    """
    采用多线程模式代用write_file函数, 实现对数据的切分
    """
    # 创建一个线程池
    threads = []

    # 将数据分割成32份
    total_data = len(data_to_process)
    num_threads = 32
    chunk_size = total_data // num_threads

    # 创建并启动多个线程
    for i in range(num_threads):
        start_idx = i * chunk_size
        end_idx = (i + 1) * chunk_size if i < num_threads - 1 else total_data
        thread = threading.Thread(target=write_file, args=(data_to_process[start_idx:end_idx], label))
        threads.append(thread)
        thread.start()

    # 等待所有线程完成
    for thread in threads:
        thread.join()


def write_to_file():
    """
    将最总结果保存到文件
    """
    fw_text = open(output_path + '/wav_text.scp', 'w')
    fw_scp = open(output_path + '/wav_accent.scp', 'w')
    assert len(text_data) == len(wav_data)
    for text in text_data:
        print(text, file=fw_text)
    for scp in wav_data:
        print(scp, file=fw_scp)
    logger.info('程序结束...')


def delete_file(language):
    """
    删除已经处理好的音频文件
    """
    dir_path = input_path + '/' + language + '/'
    for file in os.listdir(dir_path):
        if len(file.split('_')) == 2:
            os.remove(dir_path + file)


if __name__ == '__main__':
    """"""
    loading_languages()
    for language in languages.keys():
        mul_process(languages[language], language)
    write_to_file()
